from collections import deque # 实现双端队列

from nanovllm.config import Config
from nanovllm.engine.sequence import Sequence, SequenceStatus
from nanovllm.engine.block_manager import BlockManager


class Scheduler:
    """
    1.初始化以block为单位的kvcache的内存管理器
    2.设置两个队列，一个是等待队列，等待队列中都是还没开始运行的请求，
    一个是运行队列，其中都是至少在进行prefilling的请求，队列中其实用来管理发过来的请求。
    """
    def __init__(self, config: Config):
        self.max_num_seqs = config.max_num_seqs #设置最大并发序列数
        self.max_num_batched_tokens = config.max_num_batched_tokens #设置一批次处理的最大 token 数
        self.eos = config.eos #设置结束标记的 token ID
        # KV缓存块管理器
        self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)
        self.waiting: deque[Sequence] = deque() #等待队列，存储等待处理的序列
        self.running: deque[Sequence] = deque() #运行队列，存储正在处理的序列

    def is_finished(self):
        """
        检查所有序列是否都已完成处理。
        """
        return not self.waiting and not self.running

    def add(self, seq: Sequence):
        self.waiting.append(seq)

    def schedule(self) -> tuple[list[Sequence], bool]:
        scheduled_seqs = []
        num_seqs = 0
        num_batched_tokens = 0
        # 调度新到达的序列（预填充阶段），等待序列有值且未达到最大序列数时循环
        while self.waiting and num_seqs < self.max_num_seqs:
            # 从等待队列中取出第一个序列
            seq = self.waiting[0]
            # 检查当前序列的token数量不能超出最大 token 数或 KV 缓存块限制
            if (num_batched_tokens + len(seq) > self.max_num_batched_tokens or not
            self.block_manager.can_allocate(seq)):
                break
            # 增加已调度序列计数
            num_seqs += 1
            # 为序列分配KV缓存块
            self.block_manager.allocate(seq)
            # 总token数减去已缓存的token数
            num_batched_tokens += len(seq) - seq.num_cached_tokens
            seq.status = SequenceStatus.RUNNING
            self.waiting.popleft()
            self.running.append(seq)
            scheduled_seqs.append(seq)
        if scheduled_seqs:
            return scheduled_seqs, True
        # 解码阶段 对 running 队列中的序列增量生成 token
        while self.running and num_seqs < self.max_num_seqs:
            seq = self.running.popleft()
            # 当块管理器无法为当前序列追加token时，进入资源抢占循环
            while not self.block_manager.can_append(seq):
                if self.running:
                    # 抢占运行队列中最后一个序列的资源
                    self.preempt(self.running.pop())
                else:
                    self.preempt(seq)
                    break
            else:
                num_seqs += 1
                # 通知块管理器可能会追加token
                self.block_manager.may_append(seq)
                scheduled_seqs.append(seq)
        assert scheduled_seqs
        self.running.extendleft(reversed(scheduled_seqs))
        return scheduled_seqs, False
    
    # 将序列状态改为等待，并释放其资源
    def preempt(self, seq: Sequence):
        seq.status = SequenceStatus.WAITING
        self.block_manager.deallocate(seq)
        self.waiting.appendleft(seq)

     # 处理模型生成的token，更新序列状态
    def postprocess(self, seqs: list[Sequence], token_ids: list[int]) -> list[bool]:
        for seq, token_id in zip(seqs, token_ids):
            seq.append_token(token_id)  # 将生成的token添加到序列中
            # 检查序列是否应该结束（遇到EOS标记或达到最大token数）
            if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:
                seq.status = SequenceStatus.FINISHED
                self.block_manager.deallocate(seq)
                # 从运行队列中移除
                self.running.remove(seq)
